package com.shujia.flink.window

import org.apache.flink.api.common.eventtime.WatermarkStrategy
import org.apache.flink.api.common.serialization.SimpleStringSchema
import org.apache.flink.connector.kafka.source.KafkaSource
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer
import org.apache.flink.streaming.api.scala._
import org.apache.flink.streaming.api.scala.function.ProcessWindowFunction
import org.apache.flink.streaming.api.windowing.windows.GlobalWindow
import org.apache.flink.util.Collector

object Demo4ClassAvgAge {
  def main(args: Array[String]): Unit = {
    val env: StreamExecutionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment

    val source: KafkaSource[String] = KafkaSource
      .builder[String]
      //kafka broker集群列表
      .setBootstrapServers("master:9092")
      //执行消费的topic
      .setTopics("student")
      //指定消费者组，一条数据在一个组内只被消费一次
      .setGroupId("shujia")
      //读取最新的数据
      .setStartingOffsets(OffsetsInitializer.latest())
      .setValueOnlyDeserializer(new SimpleStringSchema())
      .build

    //使用kafka source
    val studentDS: DataStream[String] = env.fromSource(source, WatermarkStrategy.noWatermarks(), "Kafka Source")

    /**
     * 实时统计每个班级的平均年龄，每10个同学统计一次
     *
     */

    //取出班级和年龄
    val clazzAdnAge: DataStream[(String, Double)] = studentDS.map(line => {
      val split: Array[String] = line.split(",")
      val clazz: String = split(4)
      val age: Double = split(2).toDouble
      (clazz, age)
    })

    //安装班级分组
    val keyByDS: KeyedStream[(String, Double), String] = clazzAdnAge.keyBy(_._1)

    //开窗口
    val windowDS: WindowedStream[(String, Double), String, GlobalWindow] = keyByDS
      .countWindow(10)

    //使用底层api 计算平均值
    val avgAgeDS: DataStream[(String, Double)] = windowDS.process(new MyWindowFunction)

    avgAgeDS.print()

    env.execute()
  }

}

/**
 * ProcessWindowFunction: flink底层api,可以操作底层的数据，做更灵活统计
 */
class MyWindowFunction extends ProcessWindowFunction[(String, Double), (String, Double), String, GlobalWindow] {

  /**
   * process方法，每一个窗口执行一次
   *
   * @param clazz    ：keyby key的数据
   * @param context  ：上下游对象
   * @param elements ：窗口内的数据
   * @param out      ：用于将数据发送到下游
   */
  override def process(clazz: String,
                       context: Context,
                       elements: Iterable[(String, Double)],
                       out: Collector[(String, Double)]): Unit = {

    //计算平均年龄
    val ages: List[Double] = elements.map(_._2).toList
    val avgAge: Double = ages.sum / ages.length

    //将数据发送下游
    out.collect((clazz, avgAge))
  }
}
